from mitmproxy import ctx
from mitmproxy import http
import json
from datetime import datetime
import signal
import sys
import os
import subprocess
import winreg
from openpyxl import Workbook, load_workbook
from openpyxl.utils import get_column_letter
import os.path

class DouyinCapture:
    def __init__(self):
        self.target_url = "https://www.douyin.com/aweme/v1/web/aweme/post/"
        self.excel_file = "douyin_data.xlsx"
        self.current_row = 1
        self.init_excel()
        self.enable_proxy()
        print("程序已启动，按 Ctrl+C 停止")
        signal.signal(signal.SIGINT, self.signal_handler)
    
    def init_excel(self):
        """初始化Excel文件"""
        if not os.path.exists(self.excel_file):
            wb = Workbook()
            ws = wb.active
            headers = ['时间', 'Cookie', 'User-Agent', 'Referer', 'Play_Addr', 'Desc']
            for col, header in enumerate(headers, 1):
                ws.cell(row=1, column=col, value=header)
            wb.save(self.excel_file)
            self.current_row = 2
        else:
            wb = load_workbook(self.excel_file)
            ws = wb.active
            self.current_row = ws.max_row + 1
    
    def write_to_excel(self, data_dict):
        """写入数据到Excel"""
        try:
            wb = load_workbook(self.excel_file)
            ws = wb.active
            
            # 定义列顺序
            columns = ['时间', 'Cookie', 'User-Agent', 'Referer', 'Play_Addr', 'Desc']
            
            # 写入数据
            for col, key in enumerate(columns, 1):
                ws.cell(row=self.current_row, column=col, value=data_dict.get(key, ''))
            
            # 调整列宽
            for col in range(1, len(columns) + 1):
                ws.column_dimensions[get_column_letter(col)].width = 30
            
            wb.save(self.excel_file)
            self.current_row += 1
        except Exception as e:
            print(f"Excel写入失败: {str(e)}")

    def enable_proxy(self):
        try:
            if os.name == 'nt':
                subprocess.run('netsh winhttp set proxy 127.0.0.1:8080', shell=True)
                INTERNET_SETTINGS = winreg.OpenKey(winreg.HKEY_CURRENT_USER,
                    r'Software\Microsoft\Windows\CurrentVersion\Internet Settings',
                    0, winreg.KEY_ALL_ACCESS)
                winreg.SetValueEx(INTERNET_SETTINGS, 'ProxyEnable', 0, winreg.REG_DWORD, 1)
                winreg.SetValueEx(INTERNET_SETTINGS, 'ProxyServer', 0, winreg.REG_SZ, '127.0.0.1:8080')
        except Exception as e:
            print(f"代理设置失败: {str(e)}")
    
    def disable_proxy(self):
        try:
            if os.name == 'nt':
                subprocess.run('netsh winhttp reset proxy', shell=True)
                INTERNET_SETTINGS = winreg.OpenKey(winreg.HKEY_CURRENT_USER,
                    r'Software\Microsoft\Windows\CurrentVersion\Internet Settings',
                    0, winreg.KEY_ALL_ACCESS)
                winreg.SetValueEx(INTERNET_SETTINGS, 'ProxyEnable', 0, winreg.REG_DWORD, 0)
        except Exception as e:
            print(f"代理关闭失败: {str(e)}")
    
    def signal_handler(self, sig, frame):
        print("\n正在停止程序...")
        self.disable_proxy()
        sys.exit(0)

    def request(self, flow: http.HTTPFlow) -> None:
        if flow.request.pretty_url.startswith(self.target_url):
            # 保存请求头信息，等待响应时一起处理
            flow.request_headers = {
                'Cookie': flow.request.headers.get('cookie', ''),
                'User-Agent': flow.request.headers.get('user-agent', ''),
                'Referer': flow.request.headers.get('referer', '')
            }

    def response(self, flow: http.HTTPFlow) -> None:
        if flow.request.pretty_url.startswith(self.target_url):
            try:
                response_data = json.loads(flow.response.content.decode('utf-8'))
                aweme_list = response_data.get('aweme_list', [])
                
                for aweme in aweme_list:
                    # 获取视频地址和描述
                    video_url = aweme.get("video", {}).get("play_addr", {}).get("url_list", [-1])[-1]
                    desc = aweme.get("desc", "")
                    
                    # 准备写入Excel的数据
                    data = {
                        '时间': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                        'Cookie': flow.request_headers['Cookie'],
                        'User-Agent': flow.request_headers['User-Agent'],
                        'Referer': flow.request_headers['Referer'],
                        'Play_Addr': video_url,
                        'Desc': desc
                    }
                    
                    # 写入Excel
                    self.write_to_excel(data)
                    print(f"已记录视频: {desc[:30]}...")
                    
            except Exception as e:
                print(f"处理响应数据失败: {str(e)}")

addons = [DouyinCapture()]